"""
基于实时轨迹感知的右转车辆排队长度预测系统 - 主程序
"""

import cv2
import numpy as np
import time
import logging
from config import SystemConfig
from vehicle_detector import YOLO_OBB_Detector
from vehicle_tracker import BotSORT_Tracker
from intent_recognizer import RightTurnIntentRecognizer
from queue_analyzer import QueueAnalyzer
from length_calculator import CurvilinearLengthCalculator
from predictor import SpatioTemporalPredictor
from utils import setup_logging, draw_visualization

class QueueLengthPredictionSystem:
    """排队长度预测系统主类"""
    
    def __init__(self, config_path="config.yaml"):
        # 初始化配置
        self.config = SystemConfig(config_path)
        self.setup_logging()
        
        # 初始化各模块
        self.detector = YOLO_OBB_Detector(self.config.detection)
        self.tracker = BotSORT_Tracker(self.config.tracking)
        self.intent_recognizer = RightTurnIntentRecognizer(self.config.intent)
        self.queue_analyzer = QueueAnalyzer(self.config.queue)
        self.length_calculator = CurvilinearLengthCalculator(self.config.geometry)
        self.predictor = SpatioTemporalPredictor(self.config.prediction)
        
        # 系统状态
        self.frame_count = 0
        self.processing_times = []
        self.system_ready = True
        
    def setup_logging(self):
        """设置日志系统"""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler('system.log'),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)
        
    def process_video_stream(self, video_source=0):
        """处理视频流"""
        self.logger.info("开始处理视频流")
        
        cap = cv2.VideoCapture(video_source)
        if not cap.isOpened():
            self.logger.error("无法打开视频源")
            return
            
        while self.system_ready:
            start_time = time.time()
            
            # 读取帧
            ret, frame = cap.read()
            if not ret:
                self.logger.warning("无法读取视频帧")
                break
                
            # 处理当前帧
            result = self.process_frame(frame, self.frame_count)
            
            # 计算处理时间
            processing_time = time.time() - start_time
            self.processing_times.append(processing_time)
            
            # 可视化结果
            visualization = draw_visualization(frame, result)
            
            # 显示结果
            cv2.imshow('Queue Length Prediction', visualization)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
                
            self.frame_count += 1
            
            # 性能监控
            if self.frame_count % 100 == 0:
                self.log_performance()
                
        cap.release()
        cv2.destroyAllWindows()
        
    def process_frame(self, frame, frame_id):
        """处理单帧图像"""
        try:
            # 1. 车辆检测
            detections = self.detector.detect(frame)
            
            # 2. 多目标跟踪
            tracks = self.tracker.update(detections, frame)
            
            # 3. 右转意图识别
            right_turn_vehicles = self.intent_recognizer.identify(tracks)
            
            # 4. 排队状态分析
            queue_clusters, queue_states = self.queue_analyzer.analyze(right_turn_vehicles)
            
            # 5. 排队长度计算
            current_lengths = []
            for cluster in queue_clusters:
                length = self.length_calculator.calculate_length(cluster)
                current_lengths.append(length)
                
            current_total_length = max(current_lengths) if current_lengths else 0
            
            # 6. 未来排队预测
            prediction = self.predictor.predict(current_total_length)
            
            return {
                'frame_id': frame_id,
                'timestamp': time.time(),
                'detections': len(detections),
                'tracks': len(tracks),
                'right_turn_vehicles': len(right_turn_vehicles),
                'queue_clusters': len(queue_clusters),
                'current_length': current_total_length,
                'prediction': prediction,
                'queue_states': queue_states,
                'status': 'success'
            }
            
        except Exception as e:
            self.logger.error(f"帧处理错误: {str(e)}")
            return {
                'frame_id': frame_id,
                'timestamp': time.time(),
                'status': 'error',
                'error': str(e)
            }
    
    def log_performance(self):
        """记录性能指标"""
        if not self.processing_times:
            return
            
        avg_time = np.mean(self.processing_times[-100:])
        fps = 1.0 / avg_time if avg_time > 0 else 0
        
        self.logger.info(
            f"性能统计 - 帧数: {self.frame_count}, "
            f"平均处理时间: {avg_time:.3f}s, "
            f"预估FPS: {fps:.1f}"
        )
    
    def shutdown(self):
        """关闭系统"""
        self.logger.info("正在关闭系统...")
        self.system_ready = False
        
        # 保存最终性能报告
        if self.processing_times:
            avg_time = np.mean(self.processing_times)
            self.logger.info(f"最终平均处理时间: {avg_time:.3f}s")

def main():
    """主函数"""
    system = QueueLengthPredictionSystem()
    
    try:
        # 处理视频流（0为摄像头，或提供视频文件路径）
        system.process_video_stream("test_video.mp4")
    except KeyboardInterrupt:
        system.logger.info("用户中断处理")
    except Exception as e:
        system.logger.error(f"系统运行错误: {str(e)}")
    finally:
        system.shutdown()

if __name__ == "__main__":
    main()